#include "stdio.h"
#include "iostream"
#include "include/sampling.h"
#include "cuda.h"

void add_cuda(at::Tensor & a, torch::Tensor &b, torch::Tensor &c)
{
    //    torch::Tensor a= torch::tensor(3.0);
    //    torch::Tensor b=torch::tensor(1.0);
    //    torch::Tensor c= torch::tensor(0.0);
    const float *aa = a.data<float>();
    const float *bb = b.data<float>();
    float *cc = c.data<float>();

    int x_len = a.size(0), y_len = a.size(1);
    int size[2] = {x_len, y_len};
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
     AddLaunch(aa,bb,cc,size,stream);
}

void test_the_cuda_add_fun(){

//    float arr_1[2][3] = {{3.0, 4.0, 5.0},{1.2, 3.6, 7.8}};
//    float arr_2[2][3] = {{1.0, 1.0, 1.0},{2.,1.0,1.}};
    float arr_1[6] = {3.0, 4.0, 5.0,1.2, 3.6, 7.8};
    float arr_2[6] = {1.0, 1.0, 1.0,2.,1.0,1.};

//    nc::NdArray<float> array = {{1,2,3},{34,1,2}};
    torch::Tensor tensor_1 = torch::from_blob(arr_1,{2,3}).cuda();//rray to tensor
    std::cout<<tensor_1<<std::endl;
    torch::Tensor tensor_2 = torch::from_blob(arr_2,{2,3}).to("cuda");
    std::cout<<tensor_2<<std::endl;
    torch::Tensor  out = torch::zeros({2,3},torch::TensorOptions().dtype(torch::kFloat32)).cuda();
    
        //get the shape x
    std::cout<<"inital value"<<out<<std::endl;
    add_cuda(tensor_1,tensor_2,out);
    std::cout<<"1+2"<<out<<std::endl;

}

int main()
{
    std::cout<<"aasjchasc"<<std::endl;
    float arr_1[2][3] = {3.0, 4.0, 5.0,1.2, 3.6, 7.8};
//    torch::Tensor tensor_1 = torch::from_blob(arr_1,{2,3}).to(torch::kCUDA);//rray to tensor
    std::cout<<"ok"<<std::endl;
    test_the_cuda_add_fun();
}
    